﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Globalization;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;

namespace DeSleeper.Library.PowerManagement
{
	internal class VistaPowerHelper
	{
		private const int ERROR_MORE_DATA = 234;
		private const int ERROR_NO_MORE_ITEMS = 259;
		internal delegate uint ReadDataFunc(IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr, IntPtr buffer, ref uint bufferSize);
		internal delegate uint WriteDataFunc(IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr, IntPtr buffer, uint bufferSize);


		internal static Guid GetActiveScheme()
		{
			IntPtr buffer = IntPtr.Zero;
			try
			{

				var result = _powerGetActiveScheme(IntPtr.Zero, ref buffer);
				throwIfWin32Error(result, "Error retrieving currently active power scheme.");
				return (Guid)Marshal.PtrToStructure(buffer, typeof(Guid));
			}
			finally
			{
				Marshal.FreeHGlobal(buffer);
			}
		}

        internal static void SetActiveScheme(Guid guid)
        {
            IntPtr buffer = Marshal.AllocHGlobal(Marshal.SizeOf(guid));
            try
            {
                Marshal.StructureToPtr(guid, buffer, false);
                var result = _powerSetActiveScheme(IntPtr.Zero, buffer);
                throwIfWin32Error(result, "Error retrieving currently active power scheme.");
            }
            finally
            {
                Marshal.FreeHGlobal(buffer);
            }
        }

        internal static Guid DuplicateScheme(Guid sourceGuid)
        {
            IntPtr sourceBuffer = Marshal.AllocHGlobal(Marshal.SizeOf(sourceGuid));
            IntPtr destinationBuffer = IntPtr.Zero;
            try
            {
                Marshal.StructureToPtr(sourceGuid, sourceBuffer, false);
                var result = _powerDuplicateScheme(IntPtr.Zero, sourceBuffer, ref destinationBuffer);
                throwIfWin32Error(result, "Error retrieving currently active power scheme.");
                return (Guid)Marshal.PtrToStructure(destinationBuffer, typeof(Guid));
            }
            finally
            {
                Marshal.FreeHGlobal(sourceBuffer);
                Marshal.FreeHGlobal(destinationBuffer);
            }
        }

		internal static IEnumerable<T> GetPowerEnumeration<T>(Guid schemeGuid, Guid subgroupGuid, POWER_DATA_ACCESSOR accessType)
		{
			var index = 0;
			while (true)
			{
				T value;
				if (getPowerEnumerationValue(schemeGuid, subgroupGuid, accessType, index++, out value))
					yield return value;
				else yield break;
			}
		}

		private static bool getPowerEnumerationValue<T>(Guid schemeGuid, Guid subgroupGuid, POWER_DATA_ACCESSOR accessType, int index, out T value)
		{
			return TryReadPowerValue(
				schemeGuid, subgroupGuid, Guid.Empty,
				(IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
				 IntPtr buffer, ref uint bufferSize) =>
				_powerEnumerate(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, accessType, (uint)index,
											buffer, ref bufferSize),
				out value, 0);
		}

		internal static IEnumerable<T> GetPossibleValues<T>(Guid subgroupGuid, Guid settingGuid)
		{
			var index = 0;
			while (true)
			{
				T value;
				if (getPossibleValue(subgroupGuid, settingGuid, index++, out value))
					yield return value;
				else yield break;
			}
		}
		private static bool getPossibleValue<T>(Guid subgroupGuid, Guid settingGuid, int index, out T value)
		{
			return TryReadPowerValue(
				Guid.Empty, subgroupGuid, settingGuid,
				(IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
				 IntPtr buffer, ref uint bufferSize) =>
				_powerReadPossibleValue(rootPowerKey, subgroupGuidPtr, settingGuidPtr, IntPtr.Zero, (uint)index,
											buffer, ref bufferSize),
				out value, 0);
		}

		#region Read
		/// <exception cref="ArgumentException">Setting {0} was not found for scheme {1}..</exception>
		internal static T ReadPowerValue<T>(Guid schemeGuid, Guid subgroupGuid, Guid settingGuid, PowerValueType valueType)
		{
			T value;
			if (!TryReadPowerValue(schemeGuid, subgroupGuid, settingGuid, valueType, out value))
				throw new ArgumentException(string.Format(CultureInfo.CurrentUICulture,
					"Setting {0} was not found for scheme {1}.",
					settingGuid, schemeGuid));

			return value;
		}

		/// <exception cref="ArgumentOutOfRangeException"><c>valueType</c> is out of range.</exception>
		internal static bool TryReadPowerValue<T>(Guid schemeGuid, Guid subgroupGuid, Guid settingGuid, PowerValueType valueType, out T value)
		{
			ReadDataFunc func;
			// Zero size will query the necessary size.
			uint constantBufferSize = 0;
			switch (valueType)
			{
				case PowerValueType.AC:
					func = (IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
							IntPtr buffer, ref uint bufferSize) =>
						   _powerReadACValue(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, IntPtr.Zero,
											buffer, ref bufferSize);
					break;
				case PowerValueType.ACIndex:
					constantBufferSize = 4;
					func = (IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
							IntPtr buffer, ref uint bufferSize) =>
						   _powerReadACValueIndex(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer);
					break;
				case PowerValueType.DC:
					func = (IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
							IntPtr buffer, ref uint bufferSize) =>
						   _powerReadDCValue(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, IntPtr.Zero,
											buffer, ref bufferSize);
					break;
				case PowerValueType.DCIndex:
					constantBufferSize = 4;
					func = (IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
							IntPtr buffer, ref uint bufferSize) =>
						   _powerReadDCValueIndex(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer);
					break;
				case PowerValueType.FriendlyName:
					func = (IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
							IntPtr buffer, ref uint bufferSize) =>
						   _powerReadFriendlyName(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer, ref bufferSize);
					break;
				case PowerValueType.Description:
					func = (IntPtr rootPowerKey, IntPtr schemeGuidPtr, IntPtr subgroupGuidPtr, IntPtr settingGuidPtr,
							IntPtr buffer, ref uint bufferSize) =>
						   _powerReadDescription(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer, ref bufferSize);
					break;
				default:
					throw new ArgumentOutOfRangeException("valueType", valueType, "Unsupported power value type.");
			}


			return TryReadPowerValue(schemeGuid, subgroupGuid, settingGuid, func, out value, constantBufferSize);
		}


		internal static bool TryReadPowerValue<T>(Guid schemeGuid, Guid subgroupGuid, Guid settingGuid, ReadDataFunc func, out T value, uint bufferSize)
		{
			IntPtr schemeGuidPtr = IntPtr.Zero, subgroupGuidPtr = IntPtr.Zero, settingGuidPtr = IntPtr.Zero;
			try
			{
				schemeGuidPtr = guidToPtr(schemeGuid);
				subgroupGuidPtr = guidToPtr(subgroupGuid);
				settingGuidPtr = guidToPtr(settingGuid);

				if (bufferSize == 0)
				{
					var result = func(IntPtr.Zero, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, IntPtr.Zero, ref bufferSize);

                    if (result != ERROR_MORE_DATA)
                        throwIfWin32Error(result,
                                          string.Format(CultureInfo.InstalledUICulture,
                                                        "Unable to retrieve buffer size while reading power value. ({0}, {1}, {2})",
                                                        schemeGuid, subgroupGuid, settingGuid));
				}

				IntPtr buffer = Marshal.AllocHGlobal((int)bufferSize);
				try
				{
					var result = func(IntPtr.Zero, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer, ref bufferSize);

					if (result == ERROR_NO_MORE_ITEMS)
					{
						value = default(T);
						return false;
					}

					throwIfWin32Error(result, 
                        string.Format(CultureInfo.InstalledUICulture,
                                                        "Unable to retrieve value while reading power value. ({0}, {1}, {2})",
                                                        schemeGuid, subgroupGuid, settingGuid));

					if (typeof(T) == typeof(string))
						value = (T)(object)Marshal.PtrToStringUni(buffer);
					else
						value = (T)Marshal.PtrToStructure(buffer, typeof(T));

					return true;
				}
				finally
				{
					Marshal.FreeHGlobal(buffer);
				}
			}
			finally
			{
				Marshal.FreeHGlobal(schemeGuidPtr);
				Marshal.FreeHGlobal(subgroupGuidPtr);
				Marshal.FreeHGlobal(settingGuidPtr);
			}
		}
#endregion

		#region Write
		/// <exception cref="ArgumentOutOfRangeException"><c>valueType</c> is out of range.</exception>
		internal static void WritePowerValue<T>(Guid schemeGuid, Guid subgroupGuid, Guid settingGuid, PowerValueType valueType, T value)
			where T : IComparable<T>
		{
			WriteDataFunc func;
			switch (valueType)
			{
				case PowerValueType.AC:
				case PowerValueType.DC:
					uint index = 0;
					foreach (var possibleValue in GetPossibleValues<T>(subgroupGuid, settingGuid))
					{
						if (possibleValue.CompareTo(value) >= 0)
							break;
						index++;
					}
					// valueType + 1 = AC/DC Index
					WritePowerValue(schemeGuid, subgroupGuid, settingGuid, valueType + 1, index);
					return;
				case PowerValueType.ACIndex:
					func = (rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer, bufferSize) =>
						   _powerWriteACValueIndex(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, (uint)Marshal.ReadInt32(buffer));
					break;
				case PowerValueType.DCIndex:
					func = (rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer, bufferSize) =>
						   _powerWriteDCValueIndex(rootPowerKey, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, (uint)Marshal.ReadInt32(buffer));
					break;
				case PowerValueType.FriendlyName:
					func = _powerWriteFriendlyName;
					break;
				case PowerValueType.Description:
					func = _powerWriteDescription;
					break;
				default:
					throw new ArgumentOutOfRangeException("valueType", valueType, "Unsupported power value type.");
			}

			writePowerValue(schemeGuid, subgroupGuid, settingGuid, func, value);
		}

		private static void writePowerValue<T>(Guid schemeGuid, Guid subgroupGuid, Guid settingGuid, WriteDataFunc func, T value)
		{
			IntPtr schemeGuidPtr = IntPtr.Zero, subgroupGuidPtr = IntPtr.Zero, settingGuidPtr = IntPtr.Zero;
			try
			{
				schemeGuidPtr = guidToPtr(schemeGuid);
				subgroupGuidPtr = guidToPtr(subgroupGuid);
				settingGuidPtr = guidToPtr(settingGuid);

                IntPtr buffer = IntPtr.Zero;
			    var stringValue = value as string;
			    int bufferSize = stringValue == null ? Marshal.SizeOf(value) : stringValue.Length*2;
                try
                {
                    if (stringValue != null)
                    {
                        buffer = Marshal.StringToHGlobalUni(stringValue);
                    }
                    else
                    {
                        buffer = Marshal.AllocHGlobal(bufferSize);
                        Marshal.StructureToPtr(value, buffer, false);
                    }
                    var result = func(IntPtr.Zero, schemeGuidPtr, subgroupGuidPtr, settingGuidPtr, buffer,
                                          (uint) bufferSize);
                        throwIfWin32Error(result, "Unable to write power value.");
                }
                finally
                {
                    Marshal.FreeHGlobal(buffer);
                }
			}
			finally
			{
				Marshal.FreeHGlobal(schemeGuidPtr);
				Marshal.FreeHGlobal(subgroupGuidPtr);
				Marshal.FreeHGlobal(settingGuidPtr);
			}
		}
		#endregion

		private static IntPtr guidToPtr(Guid guid)
		{
			if (guid == Guid.Empty)
				return IntPtr.Zero;

			var guidPtr = Marshal.AllocHGlobal(Marshal.SizeOf(guid));
			Marshal.StructureToPtr(guid, guidPtr, false);
			return guidPtr;
		}

		/// <exception cref="InvalidOperationException"><c>InvalidOperationException</c>.</exception>
		private static void throwIfWin32Error(uint result, string message)
		{
			if (result != 0)
				throw new InvalidOperationException(message, new Win32Exception((int)result));
		}

		#region PInvoke declarations

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerEnumerate")]
		private static extern UInt32 _powerEnumerate(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			POWER_DATA_ACCESSOR accessFlags,
			uint index,
			IntPtr buffer,
			ref UInt32 bufferSize);

        [DllImport("powrprof.dll", SetLastError = true, EntryPoint = "PowerGetActiveScheme")]
		private static extern uint _powerGetActiveScheme(IntPtr rootPowerKey, ref IntPtr activeSchemeGuid);

        [DllImport("powrprof.dll", SetLastError = true, EntryPoint = "PowerSetActiveScheme")]
        private static extern uint _powerSetActiveScheme(IntPtr rootPowerKey, IntPtr schemeGuid);

        [DllImport("powrprof.dll", SetLastError = true, EntryPoint = "PowerDuplicateScheme")]
        private static extern uint _powerDuplicateScheme(IntPtr rootPowerKey, IntPtr sourceSchemeGuid, ref IntPtr destinationSchemeGuid);

		#region Read

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadACValue")]
		private static extern UInt32 _powerReadACValue(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr type,
			IntPtr buffer,
			ref UInt32 bufferSize);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadACValueIndex")]
		private static extern UInt32 _powerReadACValueIndex(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr valueIndex);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadDCValue")]
		private static extern UInt32 _powerReadDCValue(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr type,
			IntPtr buffer,
			ref UInt32 bufferSize);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadDCValueIndex")]
		private static extern UInt32 _powerReadDCValueIndex(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr valueIndex);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadFriendlyName")]
		private static extern UInt32 _powerReadFriendlyName(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr buffer,
			ref UInt32 bufferSize);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadDescription")]
		private static extern UInt32 _powerReadDescription(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr buffer,
			ref UInt32 bufferSize);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerReadPossibleValue")]
		private static extern UInt32 _powerReadPossibleValue(
			IntPtr rootPowerKey,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr type,
			uint index,
			IntPtr buffer,
			ref UInt32 bufferSize);
		#endregion

		#region Write

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerWriteACValueIndex")]
		private static extern UInt32 _powerWriteACValueIndex(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			uint valueIndex);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerWriteDCValueIndex")]
		private static extern UInt32 _powerWriteDCValueIndex(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			uint valueIndex);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerWriteDescription")]
		private static extern UInt32 _powerWriteDescription(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr buffer,
			uint bufferSize);

        [DllImport("powrprof.dll", SetLastError = true, CharSet = CharSet.Unicode, EntryPoint = "PowerWriteFriendlyName")]
		private static extern UInt32 _powerWriteFriendlyName(
			IntPtr rootPowerKey,
			IntPtr schemeGuid,
			IntPtr subGroupOfPowerSettingGuid,
			IntPtr settingGuid,
			IntPtr buffer,
			uint bufferSize);
		#endregion
		#endregion


	}

	[Flags]
	internal enum POWER_DATA_ACCESSOR
	{
		/// <summary>
		/// Check for overrides on AC power settings.
		/// </summary>
		ACCESS_AC_POWER_SETTING_INDEX = 0x0,
		/// <summary>
		/// Check for overrides on DC power settings.
		/// </summary>
		ACCESS_DC_POWER_SETTING_INDEX = 0x1,
		/// <summary>
		/// Check for restrictions on specific power schemes.
		/// </summary>
		ACCESS_SCHEME = 0x10,
		/// <summary>
		/// Used to enumerate subgroups with PowerEnumerate.
		/// </summary>
		ACCESS_SUBGROUPACCESS_SUBGROUP = 0x11,
		/// <summary>
		/// Used to enumerate individual power settings with PowerEnumerate.
		/// </summary>
		ACCESS_INDIVIDUAL_SETTING = 0x12,
		/// <summary>
		/// Check for restrictions on active power schemes.
		/// </summary>
		ACCESS_ACTIVE_SCHEME = 0x13,
		/// <summary>
		/// Check for restrictions on creating or restoring power schemes.
		/// </summary>
		ACCESS_CREATE_SCHEME = 0x14
	} ;
}